package org.ainlolcat.ainscience.statistics.handlers.fitting.functions;

import org.ainlolcat.ainscience.statistics.handlers.fitting.Peak;
import org.apache.commons.math.optimization.fitting.ParametricRealFunction;

import java.util.List;

/**
* @author ainlolcat
*         Date: 5/9/13
*/
public class ZeroBasedMultipleGaussFunction implements ParametricRealFunction {

    int[] shifts;
    List<Peak> peaks;

    public ZeroBasedMultipleGaussFunction(List<Peak> peaks) {
        this.peaks = peaks;
        shifts = new int[peaks.size()];
        shifts[0] = peaks.get(0).fixed ? 1 : 0;
        for (int i=1;i<peaks.size();i++){
            shifts[i] = shifts[i-1] +  (peaks.get(i).fixed ? 1 : 0);
        }
    }

    public double value(double x, double[] parameters) {
        double result = 0;
        for (int i = 0; i< peaks.size(); i++){
            Peak current = peaks.get(i);
            int shift = i > 0 ? shifts[i-1] :  0;
            double A = parameters[i * 3 + 0 - shift];
            double x0 = current.fixed ? current.x0 : parameters[i * 3 + 1 - shift];
            double w =  current.fixed ? parameters[i * 3 + 1 - shift] : parameters[i * 3 + 2 - shift];
            result += A/(w*Math.sqrt(Math.PI/2))*
                    Math.exp(-2 * Math.pow((x-x0), 2) / (Math.pow(w, 2)));
        }
        return result;
    }

    public double[] gradient(double x, double[] parameters) {
        double[] gradientVector = new double[parameters.length];
        for (int i = 0; i< peaks.size(); i++){
            Peak current = peaks.get(i);
            int shift = i > 0 ? shifts[i-1] :  0;
            double A = parameters[i * 3 + 0 - shift];
            double x0 = current.fixed ? current.x0 : parameters[i * 3 + 1 - shift];
            double w =  current.fixed ? parameters[i * 3 + 1 - shift] : parameters[i * 3 + 2 - shift];

            double dx = x0-x;
            double EXP_VALUE = Math.exp(-2 * Math.pow(dx, 2) / (Math.pow(w, 2)));      //exp(-2*dx^2/w^2)
            double OVER_W_PI_SQUARED = 1 /(w*Math.sqrt(Math.PI/2));                    //1/[w*sqrt(PI)]
            double A_OVER_W_PI_SQUARED = A*OVER_W_PI_SQUARED;                          //A/[w*sqrt(PI)]

            gradientVector[i * 3 + 0 - shift] =                                                                 //dF/dA
                    OVER_W_PI_SQUARED * EXP_VALUE;
            gradientVector[i * 3 + (current.fixed ? 1 : 2 ) - shift] =                                          //dF/dw
                    (
                    (-1)*A/(Math.pow(w,2)*Math.sqrt(Math.PI/2))
                    +
                    A_OVER_W_PI_SQUARED * (-2*dx*dx)*(-2/Math.pow(w,3))
                    )*
                    EXP_VALUE;
            if (!current.fixed)
                gradientVector[i * 3 + 1 - shift] =                                                                  //dF/dx0
                        A_OVER_W_PI_SQUARED
                        *
                        (-2/Math.pow(w,2))*(-2*dx)
                        *
                        EXP_VALUE;
        }
        return gradientVector;
    }
}
